import cv2
import numpy as np
import time
import torch
from djitellopy import Tello

# ------------------------------
# 定义 PID 控制器类
# ------------------------------
class PID:
    def __init__(self, kp, ki, kd, setpoint=0):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.setpoint = setpoint
        self.prev_error = 0
        self.integral = 0
        self.last_time = time.time()

    def compute(self, current_value):
        now = time.time()
        dt = now - self.last_time
        error = self.setpoint - current_value
        self.integral += error * dt
        derivative = (error - self.prev_error) / dt if dt > 0 else 0
        output = self.kp * error + self.ki * self.integral + self.kd * derivative
        self.prev_error = error
        self.last_time = now
        return output


# ------------------------------
# 定义 PID 控制器实例，用于控制水平和垂直方向
# ------------------------------
pid_x = PID(kp=0.1, ki=0.0, kd=0.05, setpoint=0)   # 水平误差
pid_y = PID(kp=0.1, ki=0.0, kd=0.05, setpoint=0)   # 垂直误差

##############################
# YOLOv5 模型初始化
##############################
# 使用 torch.hub 从本地加载你自定义训练的 YOLOv5 模型权重
model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5_circle.pt', source='local')
# conf=0.5 可根据需要调整

##############################
# Tello 初始化
##############################
tello = Tello()
tello.connect()
print("Battery:", tello.get_battery())
tello.streamon()
time.sleep(2)  # 等待视频流稳定

##############################
# 状态机与搜索策略参数
##############################
# 状态：SEARCH（主动扫描）、ALIGN（对准目标）、APPROACH（前进穿越）、CROSS（认为穿过）
state = "SEARCH"
lost_count = 0  # 连续未检测到目标帧计数
tol = 20  # 对齐容忍误差（像素）
forward_speed = 15  # 穿越阶段前进速度
lost_threshold = 3  # 连续丢失目标帧阈值

# 搜索阶段参数
search_rotation_speed = 30  # 水平旋转速度（yaw）
scan_counter = 0  # 水平扫描计数
scan_max = 12  # 例如12次旋转（每次1秒，大约360°）
target_search_height = 100  # 目标搜索高度（单位：厘米），认为低处目标更易检测
descend_speed = 20  # 如果高度过高时下降的速度

takeoff_done = False


def select_nearest_circle(results, frame_width, frame_height):
    """
    从 YOLOv5 检测结果中选择面积最大的目标框，
    假设目标为圆且面积越大代表距离越近。
    返回 [x1, y1, x2, y2] 或 None。
    """
    best_box = None
    best_area = 0
    for result in results.pred[0]:
        det = result.cpu().numpy()
        cls = int(det[5])
        if cls != 0:  # 只考虑类别为 circle（类别索引 0）
            continue
        area = (det[2] - det[0]) * (det[3] - det[1])
        if area > best_area:
            best_area = area
            best_box = det[:4].astype(int)
    return best_box


##############################
# 主循环
##############################
while True:
    # 起飞（仅一次）
    if not takeoff_done:
        print("Taking off")
        tello.takeoff()
        takeoff_done = True
        time.sleep(2)

    # 在搜索阶段，先检查高度是否过高
    current_height = tello.get_height()
    diff_height = target_search_height - current_height
    tello.send_rc_control(0, 0, diff_height, 0)
    time.sleep(1)


    flag = False
    yaw_sum = 0
    while True:
        # 获取当前帧
        frame = tello.get_frame_read().frame
        if frame is None:
            continue

        # 图像调整
        frame_resized = cv2.resize(frame, (416, 416))
        height, width = frame_resized.shape[:2]

        # 利用 YOLOv5 进行检测
        results = model(frame_resized, conf=0.5, verbose=False)
        target_box = select_nearest_circle(results, width, height)
        circle_detected = target_box is not None

        # 在图像上显示状态与检测结果
        cv2.putText(frame_resized, state, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        if circle_detected:
            x1, y1, x2, y2 = target_box
            cv2.rectangle(frame_resized, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(frame_resized, "circle", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)

        cv2.imshow("YOLOv5 Detection", frame_resized)


        # 状态机逻辑
        if state == "ALIGN":
            if circle_detected and target_box is not None:
                # 计算检测框中心
                x1, y1, x2, y2 = target_box
                box_center_x = (x1 + x2) // 2
                box_center_y = (y1 + y2) // 2

                # 计算误差（正值表示目标在右边或下边）
                error_x = box_center_x - (width // 2)
                error_y = box_center_y - (height // 2)
                print(f"ALIGN: error_x={error_x}, error_y={error_y}")

                # 使用 PID 计算调整命令
                adjust_x = pid_x.compute(error_x)
                adjust_y = pid_y.compute(error_y)
                # 限制调整值范围（例如 -30 到 30）
                adjust_x = max(min(int(adjust_x), 30), -30)
                adjust_y = max(min(int(adjust_y), 30), -30)

                # 发送控制命令
                # 水平调整：正值向右，负值向左
                # 垂直调整：根据 Tello API，通常正值上升（或根据你实际测试调整）
                tello.send_rc_control(adjust_x, 0, -adjust_y, 0)
                time.sleep(0.5)

                # 当误差足够小时切换状态
                if abs(error_x) < tol and abs(error_y) < tol:
                    print("Aligned with circle. Switching to APPROACH mode.")
                    state = "APPROACH"
                    lost_count = 0
            else:
                print("ALIGN: Circle not detected, hovering")
                tello.send_rc_control(0, 0, 0, 0)
                time.sleep(0.5)

        elif state == "APPROACH":
            if circle_detected and target_box is not None:
                print("APPROACH: Circle detected, moving forward slowly")
                tello.send_rc_control(0, forward_speed, 0, 0)
                lost_count = 0
            else:
                lost_count += 1
                print("APPROACH: Circle lost, count =", lost_count)
                if lost_count >= lost_threshold:
                    print("Assuming circle has been crossed.")
                    state = "CROSS"
                    tello.send_rc_control(0, 0, 0, 0)
                    time.sleep(0.5)
                    tello.land()
                    break
                else:
                    tello.send_rc_control(0, forward_speed, 0, 0)
            time.sleep(0.5)

        elif state == "CROSS":
            print("CROSS: Circle passed.")
            break

        if cv2.waitKey(1) & 0xFF == ord('q'):
            flag = True
            break

        if yaw_sum >= 360:
            yaw_sum = 0
            break
        tello.send_rc_control(0 ,0, 0, search_rotation_speed)
        yaw_sum += search_rotation_speed
        time.sleep(1)

    if flag:
        break
    target_search_height += 10

cv2.destroyAllWindows()
if takeoff_done:
    try:
        tello.land()
    except Exception as e:
        print("Landing error:", e)
